import type { Express, Request, Response, NextFunction } from 'express';
import express from 'express';
import { randomUUID } from 'node:crypto';

import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { isInitializeRequest, type JSONRPCRequest } from '@modelcontextprotocol/sdk/types.js';

// 假设 Logger 和 HttpServer 是自定义的工具模块。
// 请确保其路径在您的项目结构中是正确的。
import { Logger } from './utils/logger';
import { HttpServer } from '~/utils/http';

// --- 传输管理 (Transport Management) ---
// 一个专门用来管理所有传输实例生命周期的类。
class TransportManager {
  private sse = new Map<string, SSEServerTransport>();
  private streamable = new Map<string, StreamableHTTPServerTransport>();

  public addSseTransport(transport: SSEServerTransport): void {
    this.sse.set(transport.sessionId, transport);
    Logger.log(`SSE 传输已新增，会话 ID: ${transport.sessionId}`);
  }

  public getSseTransport(sessionId: string): SSEServerTransport | undefined {
    return this.sse.get(sessionId);
  }

  public removeSseTransport(sessionId: string): void {
    if (this.sse.delete(sessionId)) {
      Logger.log(`SSE 传输已移除，会话 ID: ${sessionId}`);
    }
  }

  public addStreamableTransport(transport: StreamableHTTPServerTransport): void {
    if (!transport.sessionId) {
      Logger.error('试图新增一个没有会话 ID 的可流式传输。');
      return;
    }
    this.streamable.set(transport.sessionId, transport);
    Logger.log(`可流式 HTTP 传输已新增，会话 ID: ${transport.sessionId}`);
  }

  public getStreamableTransport(sessionId: string): StreamableHTTPServerTransport | undefined {
    return this.streamable.get(sessionId);
  }

  public removeStreamableTransport(sessionId: string): void {
    if (this.streamable.delete(sessionId)) {
      Logger.log(`可流式 HTTP 传输已移除，会话 ID: ${sessionId}`);
    }
  }
}

// --- 路由处理器 (Route Handlers) ---

/**
 * 创建并设置 SSE 传输的路由处理器。
 * @param mcpServer 主 McpServer 实例。
 * @param transportManager 传输实例的管理员。
 * @returns 一个用于 SSE 的 Express 路由实例。
 */
function createSseRouter(mcpServer: McpServer, transportManager: TransportManager) {
  const router = express.Router();

  // 用于建立新 SSE 连接的端点。
  router.get('/', async (req: Request, res: Response, next: NextFunction) => {
    try {
      const transport = new SSEServerTransport('/messages', res);
      Logger.log(`新的 SSE 连接请求，会话 ID: ${transport.sessionId}`);

      transportManager.addSseTransport(transport);

      res.on('close', () => {
        transportManager.removeSseTransport(transport.sessionId);
        Logger.log(`SSE 连接已关闭，会话 ID: ${transport.sessionId}`);
      });

      await mcpServer.connect(transport);
      Logger.info(`SSE 连接已建立，会话 ID: ${transport.sessionId}`);
    } catch (error) {
      Logger.error('建立 SSE 连接失败:', error);
      // 将错误传递给全局错误处理器。
      next(error);
    }
  });

  // 用于在已建立的 SSE 连接上接收消息的端点。
  router.post('/messages', async (req: Request, res: Response) => {
    const sessionId = req.query.sessionId as string;
    if (!sessionId) {
      return res.status(400).send('错误的请求：缺少 sessionId 查询参数。');
    }

    const transport = transportManager.getSseTransport(sessionId);
    if (!transport) {
      return res.status(404).send(`找不到资源：没有对应此会话 ID 的 SSE 传输: ${sessionId}`);
    }

    Logger.log(`收到 SSE 消息，会话 ID: ${sessionId}`);
    await transport.handlePostMessage(req, res);
  });

  return router;
}

/**
 * 创建并设置可流式 HTTP 传输的路由处理器。
 * @param mcpServer 主 McpServer 实例。
 * @param transportManager 传输实例的管理员。
 * @returns 一个用于可流式 HTTP 的 Express 路由实例。
 */
function createStreamableHttpRouter(mcpServer: McpServer, transportManager: TransportManager) {
  const router = express.Router();

  // 处理可流式 HTTP 的初始化和后续消息。
  const handlePostRequest = async (req: Request, res: Response, next: NextFunction) => {
    const sessionId = req.headers['mcp-session-id'] as string | undefined;
    const requestBody = req.body as JSONRPCRequest;
    let transport: StreamableHTTPServerTransport | undefined;

    try {
      if (sessionId) {
        transport = transportManager.getStreamableTransport(sessionId);
        if (!transport) {
          return res.status(404).json({
            jsonrpc: '2.0',
            error: { code: -32001, message: `找不到会话: ${sessionId}` },
            id: requestBody.id ?? null,
          });
        }
        Logger.log(`重用可流式 HTTP 传输，会话 ID: ${sessionId}`);
      } else if (isInitializeRequest(requestBody)) {
        Logger.log('新的可流式 HTTP 初始化请求。');
        transport = new StreamableHTTPServerTransport({
          sessionIdGenerator: () => randomUUID(),
          onsessioninitialized: newSessionId => {
            transportManager.addStreamableTransport(transport!);
          },
        });

        transport.onclose = () => {
          if (transport?.sessionId) {
            transportManager.removeStreamableTransport(transport.sessionId);
          }
        };

        await mcpServer.connect(transport);
      } else {
        return res.status(400).json({
          jsonrpc: '2.0',
          error: { code: -32000, message: '错误的请求：请求必须是初始化请求或包含有效的会话 ID。' },
          id: requestBody.id ?? null,
        });
      }

      // 其余逻辑（例如进度通知）可以在此处或在 MCP 方法实现中处理。
      // 此示例保留进度通知以供演示。
      const progressInterval = setupProgressNotifications(mcpServer, requestBody);

      await transport.handleRequest(req, res, requestBody);

      if (progressInterval) {
        clearInterval(progressInterval);
      }
      Logger.log(`可流式 HTTP 请求已处理，会话 ID: ${transport.sessionId}`);
    } catch (error) {
      Logger.error('处理可流式 HTTP POST 请求时发生错误:', error);
      next(error);
    }
  };

  // 处理会话终止请求。
  const handleDeleteRequest = async (req: Request, res: Response, next: NextFunction) => {
    const sessionId = req.headers['mcp-session-id'] as string | undefined;
    if (!sessionId) {
      return res.status(400).send('错误的请求：DELETE 请求缺少 mcp-session-id 标头。');
    }

    const transport = transportManager.getStreamableTransport(sessionId);
    if (!transport) {
      return res.status(404).send(`找不到资源：没有对应此会话 ID 的可流式 HTTP 传输: ${sessionId}`);
    }

    Logger.log(`收到会话终止请求，会话 ID: ${sessionId}`);
    try {
      // 将请求转发到传输层以处理正常的终止流程。
      await transport.handleRequest(req, res);
    } catch (error) {
      Logger.error('处理会话终止时发生错误:', error);
      next(error);
    }
  };

  router.post('/', handlePostRequest);
  router.delete('/', handleDeleteRequest);

  return router;
}

/**
 * 设置一个模拟的进度通知间隔以供演示。
 * 在实际应用中，此逻辑可能属于模型方法执行的实际工作的一部分。
 */
function setupProgressNotifications(
  mcpServer: McpServer,
  requestBody: JSONRPCRequest
): NodeJS.Timeout | null {
  const progressToken = requestBody.params?._meta?.progressToken;
  if (!progressToken) {
    return null;
  }

  Logger.log(`正在为令牌设置进度通知: ${progressToken}`);
  let progress = 0;
  const interval = setInterval(() => {
    progress++;
    Logger.log(`正在发送进度通知: ${progress}`);
    mcpServer.server
      .notification({
        method: 'notifications/progress',
        params: { progress, progressToken },
      })
      .catch(err => {
        Logger.error('发送进度通知失败:', err);
      });
  }, 1000);

  return interval;
}

// --- 服务器注册 (Server Registration) ---

/**
 * 初始化所有 MCP 路由并将其注册到 Express 应用程序中。
 * @param mcpServer 用于处理连接的 McpServer 实例。
 */
export function registerServerRoutes(mcpServer: McpServer): void {
  const app = HttpServer.app as Express;
  const transportManager = new TransportManager();

  // --- 中间件 (Middleware) ---
  // 在您的路由之前使用 JSON 主体解析器至关重要。
  app.use(express.json());

  // --- 路由 (Routes) ---
  const sseRouter = createSseRouter(mcpServer, transportManager);
  app.use('/mcp-sse', sseRouter);

  const streamableHttpRouter = createStreamableHttpRouter(mcpServer, transportManager);
  app.use('/mcp', streamableHttpRouter);

  Logger.info('MCP 路由已成功注册。');
  Logger.info('可用传输: /mcp-sse (SSE), /mcp (可流式 HTTP)');

  // --- 全局错误处理器 (Global Error Handler) ---
  // 这应该是最后加入的中间件。
  app.use((err: Error, req: Request, res: Response, next: NextFunction) => {
    Logger.error('发生未处理的错误:', err);

    if (res.headersSent) {
      // 如果标头已发送，我们无法发送新的回应。
      // 委派给默认的 Express 错误处理器以关闭连接。
      return next(err);
    }

    res.status(500).json({
      jsonrpc: '2.0',
      error: {
        code: -32002,
        message: '发生内部服务器错误。',
        data: err.message, // 在生产环境中避免发送堆栈跟踪
      },
      id: (req.body as JSONRPCRequest)?.id ?? null,
    });
  });
}
